from typing import Dict

import gym
import numpy as np
import torch

from ....exp.logger import Logger
from .agent import ActorCritic
from .core import Core, PPOBuffer


class PPO(Core):
    """Clip版本的PPO算法，加入了KL估计用于早停策略"""

    def __init__(
        self,
        env: gym.Env,
        actor_critic: ActorCritic,
        steps_per_epoch: int = 4000,
        epochs: int = 50,
        gamma: float = 0.99,
        lam: float = 0.97,
        clip_ratio: float = 0.2,
        pi_lr: float = 3e-4,
        vf_lr: float = 1e-3,
        train_pi_iters: int = 80,
        train_v_iters: int = 80,
        max_ep_len: int = 1000,
        target_kl: float = 0.01,
        save_freq: int = 10,
        mini_batch: int = 500,
        use_norm_grad: bool = False,
        lr_decay: str = 'none',
    ) -> None:
        """函数初始化

        Args:
            env (gym.Env): 强化学习环境，要求满足gym环境接口。
            actor_critic (ActorCritic): AC智能体，要求满足基础接口。
            steps_per_epoch (int, optional): 智能体与环境交互的次数（状态-动作对的数量）。
            epochs (int, optional): 策略模型更新迭代的次数。
            gamma (float, optional): 衰减因子。
            lam (float, optional): GAE-Lambda函数中的超参数lambda。
            clip_ratio (float, optional): CLIP函数中的超参数，通常较小。
            pi_lr (float, optional): 策略更新的学习率。
            vf_lr (float, optional): 值函数更新的学习率。
            train_pi_iters (int, optional): 策略每一轮更新的最大次数，早停可能会使更新次数减少。
            train_v_iters (int, optional): 值函数每一轮更新的最大次数，早停可能会使更新次数减少。
            max_ep_len (int, optional): 一个回合的最大步数。
            target_kl (float, optional): 新旧策略的KL散度估计，用于早停，通常很小。
            save_freq (int, optional): 网络模型参数保存的频率。
            mini_batch (int, optional): 训练批次数据。
        """
        super().__init__(
            actor_critic,
            steps_per_epoch,
            epochs,
            clip_ratio,
            pi_lr,
            vf_lr,
            train_pi_iters,
            train_v_iters,
            max_ep_len,
            target_kl,
            save_freq,
            mini_batch,
            use_norm_grad,
            lr_decay,
        )
        self._env = env
        self._buf = PPOBuffer(steps_per_epoch, gamma, lam)

    @Logger.run_time("PPO._collect被调用")
    def _collect(self) -> Dict[str, torch.Tensor]:
        """收集数据"""
        o = self._env.reset()
        ep_ret, ep_len = 0, 0
        for t in range(self._steps_per_epoch):
            a, v, logp = self._ac.step(o, True)
            next_o, r, d, _ = self._env.step(a)
            ep_ret += r
            ep_len += 1
            self._buf.store(o, a, r, v, logp)
            o = next_o

            timeout = ep_len == self._max_ep_len
            terminal = d or timeout
            epoch_ended = t == self._steps_per_epoch - 1

            if terminal or epoch_ended:
                if timeout or epoch_ended:
                    _, v, _ = self._ac.step(o, True)
                else:
                    v = 0
                self._buf.finish_path(v)
                o = self._env.reset()
                ep_ret, ep_len = 0, 0
        return self._buf.get()

    @Logger.run_time("PPO._eval被调用")
    def _eval(self) -> float:
        eval_time = 5
        o = self._env.reset()
        ep_ret, ep_len = 0, 0
        reward_list = np.zeros(eval_time)
        for i in range(eval_time):
            for _ in range(self._max_ep_len):
                a, _, _ = self._ac.step(o, False)
                o, r, d, _ = self._env.step(a)
                ep_ret += r
                ep_len += 1
                if d:
                    break
            reward_list[i] = ep_ret / ep_len
        eval_val = reward_list.mean()
        Logger.info(f"eval val: {eval_val}")
        return eval_val
